import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.math.BigInteger;
import java.security.SecureRandom;

/*
 * To do: this class to implement the RSA algorithm
 * Authors: Dang Toan Tri and Nguyen Trieu Luat
 * University: HCM University of Technology
 * Department of Computer
 * Lecture: Cryptography and Network Security 
 * 
 */

// START CLASS RSA
public class RSA {
	private static final SecureRandom rand = new SecureRandom(SecureRandom.getSeed(20));
	// bit length of big prime
	private static int primeLength = 512;
	// modulus bit length = primeLength*2 => octet = (primeLength*2)/8
	private static int octet = primeLength/4;

	public static final String outFile = "out.encrypt";
	public static final String pFile = "p.numb";
	public static final String qFile = "q.numb";
	public static final String eFile = "e.numb";
	public static final String dFile = "d.numb";

	public static void setPrimeLength(int length) {
		primeLength = length;
		if (primeLength%4 == 0)
			octet = primeLength/4;
		else
			octet = primeLength/4 + 1;
	} 

	// CHECK BIG NUMBER FOR PRIMALITY
	/**
	 * 
	 * @param a : A random BigInteger to check n for primality
	 * @param n : A BigInteger number to be checked for primality
	 * @return : true if n "may be" prime,false if n is composite 
	 */
	private static boolean miller_rabin_pass(BigInteger a, BigInteger n) {
		// Calculate n-1
		BigInteger n_minus_one = n.subtract(BigInteger.ONE);
		BigInteger d = n_minus_one;
		// s is the index of the rightmost 1bit in this BigInteger d
		int s = d.getLowestSetBit();
		// Calculate d : (n-1) = (2^s).d
		d = d.shiftRight(s);
		// Calculate (a^d) mod n,if it's equal 1,n may be prime
		BigInteger a_to_power = modPow(a,d, n);
		if (a_to_power.equals(BigInteger.ONE)) return true;
		for (int i = 0; i < s-1; i++) {
			// if a^((2^i)*q) mod n == 1,n may be prime
			if (a_to_power.equals(n_minus_one)) return true;
			a_to_power = a_to_power.multiply(a_to_power).mod(n);
		}
		// Check one last time because the last iterate hasn't checked yet.
		if (a_to_power.equals(n_minus_one)) return true;
		return false;
	}
	
	// CHECK BIG NUMBER FOR PRIMALITY
	/**
	 * 
	 * @param n     : A BigInteger number to be checked for primality
	 * @param rand  : an instance of class Random or SecureRandom,used to generate a BigInteger number
	 * @return : true if n "may be" prime (but the probability that n is prime is very,very high),
	 *      false if n is composite
	 */
	private static boolean miller_rabin(BigInteger n) {
		// Check n for primality for 20 times
		for (int repeat = 0; repeat < 20; repeat++) {
			BigInteger a;
			do {
				// Select a random BigInteger a,use it to check n for primality
				a = new BigInteger(n.bitLength(), rand);
			} while (a.equals(BigInteger.ZERO));
			if (!miller_rabin_pass(a, n)) {
				return false;
			}
		}
		return true;
	}

	// GENERATE BIG PRIME NUMBER
	/**
	 * This method generate a BigInteger that is prime and have bitLength bits 
	 * @param bitLength : length of BigInteger prime in bits
	 * @return : a BigInteger prime number
	 */
	private static BigInteger genPrime(int bitLength) {
		BigInteger prime;
		do {
			prime = new BigInteger(bitLength,50,rand);
			// Check if prime is composite at the very start because miller_rabin test is
			// very time consuming
			if (prime.mod(BigInteger.valueOf(2)).equals(BigInteger.ZERO)) continue;
			if (prime.mod(BigInteger.valueOf(3)).equals(BigInteger.ZERO)) continue;
			if (prime.mod(BigInteger.valueOf(5)).equals(BigInteger.ZERO)) continue;
			if (prime.mod(BigInteger.valueOf(7)).equals(BigInteger.ZERO)) continue;
		} while (!miller_rabin(prime));
		return prime;
	}
	
	// GENERATE STRONG BIG PRIME NUMBER
	/**
	 * This method implement Gordon's algorithm to find a strong prime number
	 * @return A strong prime number
	 */
	private static BigInteger gordon_strong_prime() {
		// Gordon recommend start aSub2 and a at big value such as 2^12 to control
		//  final bit size of p, ( but i failed T_T) 
		final int startA = 4096;
		// Those values : pSub1,pAdd1 has bit lengths
		//  a bit less than half of p's bit length
		//  and pSub2's bit length is a bit less than pSub1's bit length
		BigInteger pSub2 = genPrime(primeLength/2-10);
		BigInteger pAdd1;
		// Loop until find out pAdd1 which is bigger than pSub2
		do
			pAdd1 = genPrime(primeLength/2);
		while (pAdd1.compareTo(pSub2) < 0);

		// Find the smallest prime pSub1 such that
		//	pSub1 = pSub2 * aSub2 + 1
		//  with aSub2 is a integer > 2 (pre-assigned to startA)
		BigInteger pSub1;
		BigInteger aSub2 = BigInteger.valueOf(startA);
		do {
			pSub1 = aSub2.multiply(pSub2).add(BigInteger.ONE);
			aSub2 = aSub2.add(BigInteger.ONE);
		}
		while (!miller_rabin(pSub1)); // test pSub1 for primality

		// Find smallest prime pZero such that
		// pZero = (pAdd1^(pSub1-1) - pSub1^(pAdd1+1)) mod (pSub1 * pAdd1)
		// use this property of modular
		//    (a-b) mod c = x
		//		x1 = a mod c
		//      x2 = b mod c
		// if (x1 >= x2) then x = x1 - x2
		// else x = x1 - x2 + c
		BigInteger pSub1MulAdd1 = pSub1.multiply(pAdd1);
		BigInteger pZeroTemp1 = modPow(pAdd1,pSub1.subtract(BigInteger.ONE),pSub1MulAdd1);
		BigInteger pZeroTemp2 = modPow(pSub1,pAdd1.subtract(BigInteger.ONE),pSub1MulAdd1);
		
		BigInteger pZero;
		if (pZeroTemp1.compareTo(pZeroTemp2) >= 0)
			pZero = pZeroTemp1.subtract(pZeroTemp2);
		else
			pZero = pZeroTemp1.subtract(pZeroTemp2).add(pSub1MulAdd1);
		
		BigInteger a = BigInteger.valueOf(startA);
		BigInteger p;
		do {
			p = pZero.add(a.multiply(pSub1MulAdd1));
			a = a.add(BigInteger.ONE);
		}
		while (!miller_rabin(p));	// Test p for primality
		return p;
	}
	
	// GENERATES BIGINTEGER OBJECTS : P,Q,E,D THEN WRITE THEM TO FILES
	/**
	 * Generate BigIntegers object that RSA Algorithm needs,those are 2 big prime number p,q
	 * calculate e,d,n from p,q,write them into files
	 * @param bitLength
	 * @throws FileNotFoundException
	 * @throws IOException
	 */
	public static void genBigInts() throws FileNotFoundException,IOException { 
		// Use Gordon's algorithm to find 2 strong prime number p & q
		BigInteger p = gordon_strong_prime();
		BigInteger q = gordon_strong_prime();
		// Re-calculate octet value
		int n = p.multiply(q).bitLength();
		octet = ( n % 8 == 0) ? (n/8) : (n/8+1);
		
		BigInteger[] keys = genKeys(p,q);
		BigInteger e = keys[0];
		BigInteger d = keys[1];

		// Write those numbers to files
		writeBigIntToFile(p,pFile);
		writeBigIntToFile(q,qFile);
		writeBigIntToFile(e,eFile);
		writeBigIntToFile(d,dFile);
	}
	
	// CALCULATE GREATEST COMMON DIVISOR OF TWO BIG INTEGERS
	/**
	 * This method calculate the greatest common divisor of 2 BigInteger using 
	 *  Euclidean's algorithm
	 * @param a : a BigInteger
	 * @param b : a BigInteger
	 * @return
	 */
	/*
    private static BigInteger gcd(BigInteger a,BigInteger b) {
        if (b == BigInteger.ZERO)
            return a;
        return gcd(b,a.mod(b));
    }*/
	
	// CALCULATE MULTIPLICATIVE INVERSE OF A BIG INTEGER
	/**
	 * This method calculates multiplicative inverse of d(modulo n),which is k :
	 * d.k = 1 mod n,
	 * k.d + l.n = gcd(d,n),
	 * gcd(d,n) = 1 => dk = 1 mod n,
	 * [ gcd(d,n) = rkl[0] ] 
	 * @param d : BigInteger number that you want to find its multiplicative inverse 
	 * @param n : modulus
	 * @return : multiplicative inverse of d
	 */
	private static BigInteger mulInv(BigInteger d,BigInteger n) throws ArithmeticException {
		BigInteger[] rkl = extendedEuclid (d, n);
		if (!rkl[0].equals (BigInteger.ONE)) {
			// d and n are not relatively prime. Throw an exception
			throw new ArithmeticException();
		}
		// if k > 0,return k
		// otherwise return k+n
		if (rkl[1].compareTo (BigInteger.ZERO) == 1) {
			return rkl[1];
		} else {
			return rkl[1].add(n);
		}
	}
	
	// EXTENDED EUCLIDE ALGORITHM
	/**
	 * Return an array of Big Integers [r, k, l],
	 * r = gcd(a,b) and k*d + l*n = r.
	 * Use Extended Euclid algorithm to determine these values.
	 * @param d
	 * @param n
	 * @return [r,k,l].
	 */
	private static BigInteger[] extendedEuclid (BigInteger d, BigInteger n) {
		BigInteger[] rkl = new BigInteger[3];
		if (n.equals(BigInteger.ZERO)) {
			rkl[0] = d;
			rkl[1] = BigInteger.ONE;
			rkl[2] = BigInteger.ZERO;
			return rkl;
		}
		rkl = extendedEuclid (n, d.mod (n));
		BigInteger x = rkl[1];
		BigInteger y = rkl[2];
		rkl[1] = y;
		rkl[2] = x.subtract (y.multiply (d.divide (n)));
		return rkl;
	}
	
	// CALCULATE MOD POWER OF A BIG INTEGER ( A^B MOD N )
	/**
	 * Returns a BigInteger whose value is (a^b mod n).
	 * (This method permits negative exponents.)
	 * @param a : base
	 * @param b : exponent
	 * @param n : modulus
	 * @return : (a^b) mod n
	 */
	private static BigInteger modPow(BigInteger a,BigInteger b,BigInteger n)
	{
		BigInteger result = BigInteger.ONE;
		BigInteger exponent = b;
		BigInteger t = a;
		// Loop to the last bit of exponent
		while (!exponent.equals(BigInteger.ZERO)) {
			// if exponent & 1 equals 1,it means if the last bit of exponent is 1
			// result = (result * t) mod n
			if (exponent.and(BigInteger.ONE).equals(BigInteger.ONE))
				result = result.multiply(t).mod(n);
			// shift right exponent 1 bit,it means divide exponent by 2
			exponent = exponent.shiftRight(1);
			// t = (2^t) mod n
			t = t.multiply(t).mod(n);       
		}

		if (result.compareTo(BigInteger.ZERO) < 0){
			return result.negate();
		}
		return result;
	}

	// GENERATE KEY PAIR OF RSA ALGORITHM
	/**
	 * This method calculates two pairs of keys : public key and private key base on
	 *  two big integer numbers p and q
	 *  It returns an array of 3 BigInteger [e,d,n] in which :
	 *      public key : [e,n]
	 *      private key : [d,n]
	 * @param p : a BigInteger prime
	 * @param q : another BigInteger prime
	 * @return an array of 3 BigIntegers [e,d,n] which are used as keys for RSA algorithms
	 */
	private static BigInteger[] genKeys(BigInteger p,BigInteger q) {
		BigInteger[] keys = new BigInteger[3];
		// Calculate n = p*q
		BigInteger n = p.multiply(q);
		//BigInteger phi = p.subtract(BigInteger.ONE).multiply(q.subtract(BigInteger.ONE));
		// generate big prime number e such that e is less than p and gcd(e,p) = 1
		BigInteger e;
		/*
        do
            e = genPrime(bitLength/2);
        while (!((e.compareTo(BigInteger.ONE) > 0)&&(e.compareTo(p) < 0) && gcd(e, phi).equals(BigInteger.ONE)));
		 */

		// 65537 is commonly used as a public exponent in the RSA cryptosystem.
		// This value is seen as a wise compromise,
		//since it is famously known to be prime,
		//large enough to avoid the attacks to which small exponents make RSA vulnerable,
		//and due to its low Hamming weight (number of 1 bits) can be computed
		//extremely quickly on binary computers, which often support shift and increment instructions. 
		e = BigInteger.valueOf(65537);
		// calculate d from e
		BigInteger d = getDecryptKey(e,p,q);
		// assign result array
		keys[0] = e;
		keys[1] = d;
		keys[2] = n;
		return keys;
	}

	// GET PRIVATE KEY D FROM E
	/**
	 * Calculate decryption key from e,p and q
	 * @param e : encryption partial key
	 * @param p : the BigInteger prime number has been generated
	 * @param q : the BigInteger prime number has been generated
	 * @return  : decryption partial key d 
	 */
	private static BigInteger getDecryptKey(BigInteger e,BigInteger p,BigInteger q) {
		// calculate phi(n) = (p-1)*(q-1)
		BigInteger phiN = p.subtract(BigInteger.ONE).multiply(q.subtract(BigInteger.ONE));
		// calculate n
		BigInteger n = p.multiply(q);
		// d is multiplicative inverse of e mod phi(n)
		//      d.e = 1 mod phi(n)
		BigInteger d;
		do {
			d = mulInv(e, phiN);
		}
		while (d.compareTo(BigInteger.ZERO) < 0 || d.compareTo(n) == 1);
		return d;
	}

	// WRITE A BIG INTEGERS TO FILE
	/**
	 * Write a BigInteger to the specified file
	 * @param bigInt : the BigInteger to be writen into file
	 * @param fileName : the file to write the BigInteger
	 * @throws FileNotFoundException : exception thrown if the file specified is not found
	 * @throws IOException : exception thown if there is error during file reading/writing
	 */
	private static void writeBigIntToFile(BigInteger bigInt,String fileName) throws FileNotFoundException,IOException {
		/*
    	FileOutputStream fos = new FileOutputStream(new File(fileName));
        // get byte array of the BigInteger then write it to file
        fos.write(bigInt.toByteArray());
        // CLose the FileOutputStream
        fos.close();
		 */


		FileWriter fw = new FileWriter(new File(fileName));
		fw.write(bigInt.toString());
		fw.close();


		/*
    	FileWriter fw = new FileWriter(new File(fileName));
    	fw.write(Base64.encode(bigInt.toString()));
    	fw.close();
		 */
	}
	
	// READ A BIG INTEGER FROM FILE
	/**
	 * Read a BigInteger from file
	 * @param fileName : the file contains the BigInteger
	 * @param size : size in byte of the BigInteger
	 * @return : the BigInteger that has been read
	 * @throws FileNotFoundException exception thrown if the specified file is not found
	 * @throws IOException : exception thown if there is error during file reading/writing
	 */
	private static BigInteger readBigIntFromFile(String fileName) throws FileNotFoundException,IOException {
		/*
    	FileInputStream fis = new FileInputStream(new File(fileName));
        byte[] b = new byte[size];
        fis.read(b);
        BigInteger bigInt = new BigInteger(b);
        fis.close();
        return bigInt;*/


		File inFile = new File(fileName);
		FileInputStream fis = new FileInputStream(inFile);
		byte[] b = new byte[(int)inFile.length()];
		fis.read(b);
		String s = new String(b);
		BigInteger bigInt = new BigInteger(s);
		fis.close();
		return bigInt;

		/*
    	File inFile = new File(fileName);
    	FileInputStream fis = new FileInputStream(inFile);
        byte[] b = new byte[(int)inFile.length()];
        fis.read(b);
        String s = new String(b);
        BigInteger bigInt = new BigInteger(Base64.decode(s));
        fis.close();
        return bigInt;
		 */
	}

	// INTERNAL METHOD - ENCRYPT A FILE USING RSA ALGORITHM
	/**
	 * Encrypt the specified inFile file using recipient's public key pair (e,n),
	 *  store the cipher into outFile  
	 * @param inFile : file that store the plain message
	 * @param outFile : file that store the cipher text
	 * @param e : exponent of public key
	 * @param n : modulus of public key
	 * @throws FileNotFoundException
	 * @throws IOException
	 */
	private static void encryptRSA(String inFile,String outFile,BigInteger e,BigInteger n) throws FileNotFoundException,IOException {
		FileInputStream fis = new FileInputStream(new File(inFile));
		FileOutputStream fos = new FileOutputStream(new File(outFile));
		BigInteger bigInt = null;
		BigInteger enBig = null;
		// Maximum length of plain message block is octet-11,but for sure,we should
		// just encrypt a block of octet-11 one time
		byte[] b = new byte[octet-12];
		int i = fis.read(b);
		while (i > 0) {
			if ( i < octet-11) {
				byte[] lastBytes = new byte[i];
				for (int j = 0; j < i ; j++)
					lastBytes[j] = b[j];
				bigInt = new BigInteger(lastBytes);
			}
			else
				// Create BigInteger object from the byte array read
				bigInt = new BigInteger(b);
			// Encrypt the BigInteger
			enBig = modPow(bigInt,e,n);
			// Write the cipher into file
			byte[] enBigByteArr = enBig.toByteArray();
			int enBigSize = enBigByteArr.length;

			if ( enBigSize < (octet+1)) {
				byte[] tempArr = new byte[octet+1];
				int delta = octet +1 - enBigSize;
				for (int j = 0 ; j < delta ; j++) {
					tempArr[j] = 0;
				}
				for (int j = delta ; j < (octet+1) ; j++) {
					tempArr[j] = enBigByteArr[j-delta];
				}
				fos.write(tempArr);
			}
			else
				fos.write(enBigByteArr);
			// Read another block of data
			i = fis.read(b);
		}
		// Close all connections
		fis.close();
		fos.close();

	}
	
	// INTERNAL METHOD - DECRYPT A FILE USING RSA ALGORITHM
	private static void decryptRSA(String inFile,String outFile,BigInteger d,BigInteger n) throws FileNotFoundException,IOException {
		FileInputStream fis = new FileInputStream(new File(inFile));
		FileOutputStream fos = new FileOutputStream(new File(outFile));
		BigInteger bigInt = null;
		BigInteger deBig = null;
		byte[] b = new byte[octet+1];
		int i = fis.read(b);
		while (i > 0) {
			if (b[0] == 0) {
				int count = 0;
				for (int j = 0 ; j < (octet+1) ; j++) {
					if (b[j] == 0) {
						if ((j < octet) && b[j+1] < 0) {
							break;
						}
						else
							count++;
					}

					else
						break;
				}
				if (count > 0) {
					byte[] newByte = new byte[octet+1-count];
					for (int j = 0 ; j < (octet+1-count) ; j++) {
						newByte[j] = b[count+j];
					}
					bigInt = new BigInteger(newByte);
				}
				else 
					bigInt = new BigInteger(b);
			}
			else
				bigInt = new BigInteger(b);
			deBig = modPow(bigInt,d,n);
			fos.write(deBig.toByteArray());
			i = fis.read(b);
		}
		fis.close();
		fos.close();

	}

	// ENCRYPT A FILE
	public static void encryptFile(String inFile) throws FileNotFoundException,IOException {
		BigInteger p = readBigIntFromFile(pFile);
		BigInteger q = readBigIntFromFile(qFile);
		BigInteger e = readBigIntFromFile(eFile);
		BigInteger n = p.multiply(q);
		encryptRSA(inFile,outFile,e,n);
	}
	
	// DECRYPT A FILE
	public static void decryptFile(String decryptFile) throws FileNotFoundException,IOException {

		BigInteger p = readBigIntFromFile(pFile);
		BigInteger q = readBigIntFromFile(qFile);
		BigInteger d = readBigIntFromFile(dFile);
		BigInteger n = p.multiply(q);
		decryptRSA(outFile,decryptFile,d,n);
	}
	
	// MAIN METHOD
	public static void main(String[] args) {
		int length = 0;
		try {
			length = Integer.parseInt(args[0]);
			if (length < 500) {
				System.err.println("You must enter a number bigger than 500");
				System.exit(1);
			}
		}
		catch (NumberFormatException nfe) {
			System.err.println("First argument must be a number bigger than 500");
			System.exit(1);
		}

		// generates p,q,e,d
		try {
			setPrimeLength(length);
			genBigInts();
		}
		catch (FileNotFoundException fnfe) {
			System.err.println("File Not Found Exception - File(s) cannot be found.");
			System.exit(1);
		}
		catch (IOException e) {
			System.err.println("IO Exception.");
			System.exit(1);
		}

		// Encrypt file
		try {
			encryptFile(args[1]);
		}
		catch (FileNotFoundException fnfe) {
			System.err.println("File Not Found Exception - File " + args[1]+ " cannot be found.");
			System.exit(1);
		}
		catch (IOException e) {
			System.err.println("IO Exception.");
			System.exit(1);
		}

		// Decrypt file
		try {
			decryptFile("outputDecrypt.txt");
		}
		catch (FileNotFoundException fnfe) {
			System.err.println("File Not Found Exception - File " + args[1]+ " cannot be found.");
			System.exit(1);
		}
		catch (IOException e) {
			System.err.println("IO Exception.");
			System.exit(1);
		}
	}

}
// END CLASS RSA




// START CLASS Base64
class Base64 {
	// Base64 code
	private static final String base64code = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
			+ "abcdefghijklmnopqrstuvwxyz" + "0123456789" + "+/";
	
	// Pad zero bytes 
	private static byte[] zeroPad(int length, byte[] bytes) {
		byte[] padded = new byte[length]; // initialized to zero by JVM
		System.arraycopy(bytes, 0, padded, 0, bytes.length);
		return padded;
	}
	
	// Encode a string using Base64 format
	public static String encode(String string) {
		
		String encoded = "";
		byte[] stringArray;
		try {
			stringArray = string.getBytes("UTF-8");  // use appropriate encoding string!
		} catch (Exception ignored) {
			stringArray = string.getBytes();  // use locale default rather than croak
		}
		// determine how many padding bytes to add to the output
		int paddingCount = (3 - (stringArray.length % 3)) % 3;
		// add any necessary padding to the input
		stringArray = zeroPad(stringArray.length + paddingCount, stringArray);
		// process 3 bytes at a time, churning out 4 output bytes
		// worry about CRLF insertions later
		for (int i = 0; i < stringArray.length; i += 3) {
			int j = ((stringArray[i] & 0xff) << 16) +
					((stringArray[i + 1] & 0xff) << 8) + 
					(stringArray[i + 2] & 0xff);
			encoded = encoded + base64code.charAt((j >> 18) & 0x3f) +
					base64code.charAt((j >> 12) & 0x3f) +
					base64code.charAt((j >> 6) & 0x3f) +
					base64code.charAt(j & 0x3f);
		}
		
		return encoded + "==".substring(0, paddingCount);
	}

	public static String decode(String string) {
		int length = string.length();
		int paddingCount;
		if(string.charAt(length-2) == '=')
			paddingCount = 2;
		else if (string.charAt(length-1) == '=')
			paddingCount = 1;
		else paddingCount = 0;

		string = string.substring(0,string.length()-paddingCount);
		byte[] decodedArray = new byte[string.length()*3/4];
		int j;
		int byteIndex = 0;
		for (int i = 0 ; i < string.length() ; i += 4) {		
			j = base64code.indexOf(string.charAt(i));
			j <<= 6;
			j |= base64code.indexOf(string.charAt(i+1));
			j <<= 6;
			j |= base64code.indexOf(string.charAt(i+2));
			j <<= 6;
			j |= base64code.indexOf(string.charAt(i+3));

			decodedArray[byteIndex+2] = (byte)(j);
			j >>= 8;
			decodedArray[byteIndex+1] = (byte)(j);
			j >>= 8;
			decodedArray[byteIndex] = (byte)(j);
			j >>= 8;
			byteIndex += 3;
		}

		byte[] returnArr = new byte[decodedArray.length - paddingCount];
		System.arraycopy(decodedArray,0,returnArr,0,decodedArray.length - paddingCount);
		return new String(returnArr);
	}
}
// END CLASS Base64



// http://people.csail.mit.edu/rivest/pubs/RS01.version-1999-11-22.pdf
// $ javac e://eclipse/workspace/RSA/src/RSA.java
// $ java -cp bin RSA 512 input.txt   